import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import numpy as np

import os
import pickle
import time



from data import *
from model.refinedet_vgg import build_refinedet
from data.coco import COCODetectionDataset, COCO_CLASSES

DEVICE = torch.device("cuda")
TRAINED_MODEL = '/media/slytheringe/Data/develop/Deep Learning/ref-det-pytorch1.5/weights/RefineDet512_COCO_119999.pth'
SAVE_FOLDER = './results'

class Timer(object):
    """A simple timer."""
    def __init__(self):
        self.total_time = 0.
        self.calls = 0
        self.start_time = 0.
        self.diff = 0.
        self.average_time = 0.

    def tic(self):
        # using time.time instead of time.clock because time time.clock
        # does not normalize for multithreading
        self.start_time = time.time()

    def toc(self, average=True):
        self.diff = time.time() - self.start_time
        self.total_time += self.diff
        self.calls += 1
        self.average_time = self.total_time / self.calls
        if average:
            return self.average_time
        else:
            return self.diff

def test_net(save_folder, net, dataset, num_classes, thresh=0.05):
    # save_folder: the folder where to save detection results
    # num_classes: contains background class
    if not os.path.exists(save_folder):
        os.mkdir(save_folder)
    num_images = len(dataset)
    # all detections are collected into:
    #    all_boxes[cls][image] = N x 5 array of detections in
    #    (x1, y1, x2, y2, score)
    all_boxes = [[[] for _ in range(num_images)]
                    for _ in range(num_classes)]
    _t = {'im_detect': Timer(), 'misc': Timer()}
    det_file = os.path.join(save_folder, 'detections.pkl')
    with torch.no_grad():
        for i in range(num_images):
            im, gt, h, w = dataset.pull_item(i)

            x = im.unsqueeze(0).to(DEVICE)
            _t['im_detect'].tic()
            detections = net(x).detach()
            detect_time = _t['im_detect'].toc(average=False)
            _t['misc'].tic()
            # skip j = 0, because it's the background class
            for j in range(1, detections.size(1)):
                dets = detections[0, j, :]  
                # make sure confidence is greater than 0
                mask = dets[:, 0].gt(0.).expand(5, dets.size(0)).t()
                '''
                ???此处的数据形式不是很清楚，为什么要转置
                '''
                dets = torch.masked_select(dets, mask).view(-1, 5)
                if dets.size(0) == 0:
                    continue
                boxes = dets[:, 1:]
                boxes[:, 0] *= w
                boxes[:, 2] *= w
                boxes[:, 1] *= h
                boxes[:, 3] *= h
                scores = dets[:, 0].cpu().numpy()
                '''
                ???这边干啥了
                '''
                cls_dets = np.hstack((boxes.cpu().numpy(), scores[:, np.newaxis])).astype(np.float32,copy=False)
                all_boxes[j][i] = cls_dets
            eval_time = _t['misc'].toc()
            print('im_detect:{:d}/{:d} detection_time:{:.3f}s eval_time:{:.3f}'.format(i + 1, num_images, detect_time, eval_time))

    with open(det_file, 'wb') as f:
        pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)

    print('Evaluating detections')

    dataset.evaluate_detections(all_boxes, save_folder)

        


if __name__ == '__main__':
    # load net
    num_classes = len(COCO_CLASSES) + 1 # +1 for background
    net = build_refinedet('test', num_classes)
    net.load_state_dict(torch.load(TRAINED_MODEL))
    net.eval()
    # load data
    dataset = COCODetectionDataset('./data/coco', 'minival2014',
                                     transform=BaseTransform(512, MEANS))
    net = net.to(DEVICE)
    cudnn.benchmark = True
    # evaluation
    test_net(SAVE_FOLDER, net, dataset, num_classes, 0.01)

    